import sys
from pathlib import Path
from ament_index_python.packages import get_package_share_directory
PATH_ROOT = Path(get_package_share_directory('manipulation')).parents[3]
PATH_SRC = PATH_ROOT / "src"
sys.path.append(str(PATH_SRC))

import cv2
import numpy as np
from manipulation.scripts.detection.cameras import D435Camera as Camera
from apriltag import apriltag

special_tag_size = {
}

def get_extrinsics_matrix_pnp(xs_img: np.ndarray, xs_world: np.ndarray, K: np.ndarray, dist=None):
    assert len(xs_img) == len(xs_world) and len(xs_img) >= 4
    if len(xs_img) > 4:
        xs_img = xs_img[:4]
        xs_world = xs_world[:4]
    _, R_vec, T = cv2.solvePnP(xs_world, xs_img, K, distCoeffs=dist, flags=cv2.SOLVEPNP_ITERATIVE)
    R, _ = cv2.Rodrigues(R_vec)
    return R, T[:, 0]

class AprilTagDetection:
    corner_world_ratio_pos = [(-0.5, -0.5, 0), (0.5, -0.5, 0), (0.5, 0.5, 0), (-0.5, 0.5, 0)]

    def __init__(self, tag_family="tag36h11"):
        self.detector = apriltag(tag_family)

    def detect(self, rgb_img, K=None, dist=None, tag_size=0.04, verbose=False):
        gray = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2GRAY)
        tags = self.detector.detect(gray)

        tag_infos: dict[int, list[dict]] = dict()
        for tag in tags:
            corners = np.array(tag['lb-rb-rt-lt'], dtype=np.float32)  # [4,2]
            xs_img = []
            xs_world = []

            for idx in range(4):
                x_img = corners[idx]
                ts = tag_size if tag['id'] not in special_tag_size else special_tag_size[tag['id']]
                x_world = np.array(self.corner_world_ratio_pos[idx], np.float32) * ts
                xs_img.append(x_img)
                xs_world.append(x_world)

            xs_img = np.array(xs_img, np.float32)
            xs_world = np.array(xs_world, np.float32)

            R, t = None, None
            if K is not None:
                R, t = get_extrinsics_matrix_pnp(xs_img, xs_world, K, dist)

            id = int(tag['id'])
            info = {
                'id': id,
                'corners': corners,
                'R': R,
                't': t,
            }
            if id not in tag_infos:
                tag_infos[id] = [info]
            else:
                tag_infos[id].append(info)
                if verbose:
                    print(f"[WARNING]: Exist {len(tag_infos[id])}-multiple tag_{id=}")
        return tag_infos

    def draw_detected_corners(self, rgb_img, K=None, dist=None, tag_size=0.04, verbose=False, draw_text=True):
        img = rgb_img.copy()
        tag_infos = self.detect(img, K, dist, tag_size, verbose)

        for id in tag_infos:
            for info in tag_infos[id]:
                colors = [(255,0,0), (0,255,0), (0,0,255), (0,255,255)]
                for idx, corner in enumerate(info['corners']):
                    corner = corner.astype(np.int32)
                    cv2.circle(img, corner, 5, colors[idx], -1)
                text_org = np.mean(info['corners'], axis=0).astype(np.int32)
                if draw_text:
                    cv2.putText(img, f"Tag{id}", text_org, cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255), thickness=2)

        return img, tag_infos

def debug_draw_corner_detect():
    april_tag_detection = AprilTagDetection()
    camera = Camera()
    while True:
        rgb_img, depth_img = camera.get_frame()
        corners_img, _ = april_tag_detection.draw_detected_corners(rgb_img)
        cv2.imshow("debug corners", corners_img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

def debug_detect_and_draw_axis():
    axis_colors = [(0,0,255), (0,255,0), (255,0,0)]
    april_tag_detection = AprilTagDetection()
    camera = Camera()
    while True:
        rgb_img, depth_img = camera.get_frame()
        img, tag_infos = april_tag_detection.draw_detected_corners(
            rgb_img=rgb_img,
            K=camera.K,
            dist=np.array(camera.dist, np.float32)
        )
        for id in tag_infos:
            for info in tag_infos[id]:
                axis_world_pos = np.array([[0,0,0], [0.1,0,0], [0,0.1,0], [0,0,0.1]], np.float32)
                axis_img_pos = camera.world_pos2img_pos(axis_world_pos, info['R'], info['t'])
                origin_img_pos = axis_img_pos[0]
                for idx, color in enumerate(axis_colors):
                    img = cv2.line(img, origin_img_pos, axis_img_pos[idx+1], color, thickness=2)
                img = cv2.circle(img, origin_img_pos, radius=5, color=(255,255,255), thickness=-1)
        cv2.imshow("debug draw axis", img)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

if __name__ == '__main__':
    # debug_draw_corner_detect()
    debug_detect_and_draw_axis()
